import networkx as nx
import json
from pydantic import BaseModel
from typing import List
from .types import ParameterMapping

class OperatorComponent(BaseModel):
    """
    OperatorComponent is a model that represents a component in the algorithm graph.

    Attributes:
        component_name (str): The name of the component. (Not used yet)
        component_label (str): The label of the component.
        component_id (str): The unique identifier of the component.
        component_graph_id (str): The unique identifier of the component in the graph.

    Methods:
        from_component(cls, id, component):
            Creates an instance of OperatorComponent from a given component dictionary.
    """
    component_name: str #not use yet
    component_label: str
    component_id: str
    component_graph_id: str # use graph id to unique identify the component

    @classmethod
    def from_component(cls, id, component):
        """
        Creates an instance of OperatorComponent from a given component dictionary.

        Args:
            id (str): The unique identifier of the component in the graph.
            component (dict): The dictionary containing component details.

        Returns:
            OperatorComponent: An instance of OperatorComponent.
        """
        return cls(component_name = "NOT USED YET", component_label = component['metadata']['label'], component_id = component['component'], component_graph_id = id)


class AlgoGraph:
    """
    A class representing an algorithm graph.

    This class creates a directed acyclic graph (DAG) from a JSON schema,
    with specified start and stop nodes. It provides methods for traversing
    the graph and retrieving information about node connections.

    Attributes:
        start_node (str): The identifier of the start node.
        stop_node (str): The identifier of the stop node.
        node_id_mapping_graph_id (dict): A mapping of component IDs to graph IDs.
        graph (nx.DiGraph): The NetworkX DiGraph object representing the algorithm.
        execution_order_ids (set): A set of node IDs in execution order.

    Args:
        json_schema (str): A JSON string representing the graph structure.
        start_node (str): The identifier of the start node.
        stop_node (str): The identifier of the stop node.

    Raises:
        Exception: If the graph is not a DAG (contains cycles).
        ValueError: If the start node has incoming edges or the stop node has outgoing edges.
    """

    def __init__(self, json_schema, start_node, stop_node):
        self.start_node = start_node
        self.stop_node = stop_node
        self.node_id_mapping_graph_id = {}
        G = nx.DiGraph()

        data = json.loads(json_schema)
        for node_graph_id, process in data['processes'].items():
            self.node_id_mapping_graph_id[process['component']] = node_graph_id

            G.add_node(node_graph_id, process=process, label=process['metadata']['label'])

        for connection in data['connections']:
            src = connection['src']['process']
            tgt = connection['tgt']['process']
            G.add_edge(src, tgt)

        if not nx.is_directed_acyclic_graph(G):
            raise Exception("It's not DAG, has cycle in it")

        if start_node in self.node_id_mapping_graph_id and G.in_degree(self.node_id_mapping_graph_id[start_node]) != 0:
            raise ValueError(f"The start node '{start_node}' must have no incoming edges.")

        if stop_node in self.node_id_mapping_graph_id and G.out_degree(self.node_id_mapping_graph_id[stop_node]) != 0:
            raise ValueError(f"The stop node '{stop_node}' must have no outgoing edges.")

        self.graph = G
        self.execution_order_ids = self.get_execution_order_ids(self.start_node)

    def find_downstream_nodes(self, node_graph_id) -> List[OperatorComponent]:
        """
        Find all downstream nodes of a given node that are in the execution order.

        Args:
            node_graph_id (str): The graph ID of the node to find downstream nodes for.

        Returns:
            List[OperatorComponent]: A list of OperatorComponent objects representing the downstream nodes.
        """
        execution_order_ids = self.execution_order_ids
        return [OperatorComponent.from_component(n, self.graph.nodes[n]['process'])
                for n in list(self.graph.successors(node_graph_id)) if n in execution_order_ids]

    def find_upstream_nodes(self, node_graph_id) -> List[OperatorComponent]:
        """
        Find all upstream nodes of a given node that are in the execution order.

        Args:
            node_graph_id (str): The graph ID of the node to find upstream nodes for.

        Returns:
            List[OperatorComponent]: A list of OperatorComponent objects representing the upstream nodes.
        """
        execution_order_ids = self.execution_order_ids
        return [OperatorComponent.from_component(n, self.graph.nodes[n]['process'])
                for n in list(self.graph.predecessors(node_graph_id)) if n in execution_order_ids]

    def get_upstream_nodes_mappings(self, node_graph_id) -> List[ParameterMapping]:
        """
        Get parameter mappings for upstream nodes of a given node.

        Args:
            node_graph_id (str): The graph ID of the node to get upstream mappings for.

        Returns:
            List[ParameterMapping]: A list of ParameterMapping objects representing the upstream node mappings.
        """
        ret: List[ParameterMapping] = []
        execution_order_ids = self.execution_order_ids
        n = node_graph_id
        if n in execution_order_ids:
            cur_node = self.graph.nodes[n]
            metadata = cur_node['process']['metadata']
            if 'parameterMapping' in metadata.keys():
                mapping_info = metadata['parameterMapping']
                for param_name, map_info in mapping_info.items():
                    p = ParameterMapping(
                        src_operator_graph_id = map_info['nodeId'],
                        source_operator_id = map_info['operatorId'],
                        dst_operator_graph_id = n,
                        dst_operator_id = cur_node['process']['component'],
                        src_param_name = map_info['name'],
                        dst_param_name = param_name
                    )
                    ret.append(p)
        return ret

    def get_all_node_mappings(self) -> List[ParameterMapping]:
        """
        Get parameter mappings for all nodes in the execution order.

        Returns:
            List[ParameterMapping]: A list of ParameterMapping objects representing all node mappings.
        """
        ret: List[ParameterMapping] = []
        for node_graph_id in self.execution_order_ids:
            ret.extend(self.get_upstream_nodes_mappings(node_graph_id))
        return ret


    def get_execution_order_ids(self, node_graph_id):
        return {comp.component_graph_id for comp in self.get_execution_order(node_graph_id)}

    #def get_execution_order(self, node_name) -> List[OperatorComponent]:
    #    execution_order_ids = list(reversed(list(nx.dfs_postorder_nodes(self.graph, source=node_name))))
    #    execution_order_ids = [n for n in execution_order_ids if n not in {self.start_node, self.stop_node}]
    #    #TODO: reemove node that node connect to stop
    #    return [OperatorComponent.from_component(n, self.graph.nodes[n]['process']) for n in execution_order_ids]

    #TODO: add test for this
    def get_execution_order(self, node_graph_id) -> List[OperatorComponent]:
        # 获取从 start_node 开始的所有节点
        start_node = self.start_node
        end_node = self.stop_node

        reachable_from_start = set(nx.descendants(self.graph, start_node)).union({start_node})
        reachable_to_end = set(nx.ancestors(self.graph, end_node)).union({end_node})

        # 取交集，只保留既能从 start_node 到达又能到达 end_node 的节点
        valid_nodes = reachable_from_start.intersection(reachable_to_end)

        # 获取从 node_name 开始的 DFS 后序遍历的节点
        execution_order_ids = list(reversed(list(nx.dfs_postorder_nodes(self.graph, source=node_graph_id))))

        # 过滤掉不在 valid_nodes 集合中的节点
        execution_order_ids = [n for n in execution_order_ids if n in valid_nodes and n not in {self.start_node, self.stop_node}]

        return [OperatorComponent.from_component(n, self.graph.nodes[n]['process']) for n in execution_order_ids]
